# Dickstein, Ho, and Mark (2023)
# This script runs a loop in which multiple counterfactuals are estimated, each with a 
# different set of counterfactual parameters are used. This analysis creates the output 
# presented in figures 5 through 7 in the paper.

# * # * # * # * # * # * #
# PRELIMINARIES         #
# * # * # * # * # * # * #
setwd("../../../library")
source("PreliminariesCode.R")

find_names <- function(name_vec, start){
  unlist(lapply(strsplit(name_vec[which(type_param_vec %in% start)], "\\."), function(x) x[3]))
}

# Are you providing your own parameters?
load_params = F

# Are the parameters being created below?
create_param_mat <- T

# Which rating areas?
ravec = 1:7

# Which years?
yrvec = 2016

# What percent of the premium is covered by an employer?
nu = .65

# What is the behavioral share
beshr <- .01

# What alpha-psi cutoff to use: 
alphampsicutoff <- .05

# What ACG cutoff to use: 
acgHL_boundary <- 1.76

# What is the base markup?
mrkup_0 <- 0

# What is psi_upperbound?
psi_upperbound <- .3

# Which spec do you want to use?
spectouse <- "lowrisk_simplemh_censor0.2"

# Where is the counterfactual folder?
counterfactual_folder <- paste0(project_folder, "/analysis/counterfactuals")

# Where is the output folder
output_folder <- paste0(project_folder, "/analysis/counterfactuals/loop_specs/output")
dir.create(output_folder, recursive = T)

# where is datasets folder
data_folder <- paste0(project_folder, "/data")

# Loading counterfactual functions 
source(paste0(counterfactual_folder, "/library/loop/DA03aRunCounterfactualFunctionsDL.R"))
source(paste0(counterfactual_folder, "/library/loop/DA03aRunCounterfactualFunctions_psiDL.R"))
source(paste0(counterfactual_folder, "/library/loop/DA03aRunCounterfactualFunctions_alphaDL.R"))

# Which counterfactuals do you want to run? 
run_which_ctfls <- c(0, 1, 5)

# Which columns do you want to create? 
create_which_columns <- c(0, 1, 5)

# ~ # ~ # V CREATING PARAM MAT HERE V # ~ # ~ #
if(create_param_mat == T){
  
  # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ #
  # Load base parameter values    #
  # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ #
  
  # Loading individual market demand params
  IndMarketEstimates_0 <- fread(paste0(project_folder, "/analysis/estimationcode/specs/", spectouse, "_main/output/solution.csv"))
  
  ## Kondo.
  IndMarketEstimates_0 <- IndMarketEstimates_0[, -c(1,2)]
  
  # Loading the number of beta1, 2, 3, 0 parameters
  IndMarketlabels_0 <- fread(paste0(project_folder, "/analysis/estimationcode/specs/", spectouse, "_main/output/covar_lens.csv"))
  IndMarSE_0 <- cumsum(IndMarketlabels_0$x)
  
  # Putting them into an easy to read data table
  NamedBetaList_0 <- list(
    beta1 = as.numeric(IndMarketEstimates_0[1, 1:IndMarSE_0[1]]),
    beta2 = as.numeric(IndMarketEstimates_0[1, (1 + IndMarSE_0[1]):IndMarSE_0[2]]),
    beta3 = as.numeric(IndMarketEstimates_0[1, (1 + IndMarSE_0[2]):IndMarSE_0[3]]),
    beta0 = as.numeric(IndMarketEstimates_0[1, (1 + IndMarSE_0[3]):IndMarSE_0[4]])
  )
  
  # Adding the original variable names to list
  names(NamedBetaList_0[[1]]) <- names(IndMarketEstimates_0)[1:IndMarSE_0[1]]
  names(NamedBetaList_0[[2]]) <- names(IndMarketEstimates_0)[(1 + IndMarSE_0[1]):IndMarSE_0[2]]
  names(NamedBetaList_0[[3]]) <- names(IndMarketEstimates_0)[(1 + IndMarSE_0[2]):IndMarSE_0[3]]
  names(NamedBetaList_0[[4]]) <- names(IndMarketEstimates_0)[(1 + IndMarSE_0[3]):IndMarSE_0[4]]
  
  # Do the same thing for the switcher beta variables
  SGMarketEstimates_0 <- fread(paste0(project_folder, "/analysis/estimationcode/specs/", spectouse, "_switcher/output/solution.csv"))
  SGMarketEstimates_0 <- SGMarketEstimates_0[, -c(1,2)]
  SGMarketlabels_0 <- fread(paste0(project_folder, "/analysis/estimationcode/specs/", spectouse, "_switcher/output/covar_lens.csv"))
  SGMarSE_0 <- cumsum(SGMarketlabels_0$x)
  SGNamedBetaList_0 <- list(
    beta1 = as.numeric(SGMarketEstimates_0[1, 1:SGMarSE_0[1]]),
    beta2 = as.numeric(SGMarketEstimates_0[1, (1 + SGMarSE_0[1]):SGMarSE_0[2]]),
    beta3 = as.numeric(SGMarketEstimates_0[1, (1 + SGMarSE_0[2]):SGMarSE_0[3]]),
    beta0 = as.numeric(SGMarketEstimates_0[1, (1 + SGMarSE_0[3]):SGMarSE_0[4]])
  )
  names(SGNamedBetaList_0[[1]]) <- names(SGMarketEstimates_0)[1:SGMarSE_0[1]]
  names(SGNamedBetaList_0[[2]]) <- names(SGMarketEstimates_0)[(1 + SGMarSE_0[1]):SGMarSE_0[2]]
  names(SGNamedBetaList_0[[3]]) <- names(SGMarketEstimates_0)[(1 + SGMarSE_0[2]):SGMarSE_0[3]]
  names(SGNamedBetaList_0[[4]]) <- names(SGMarketEstimates_0)[(1 + SGMarSE_0[3]):SGMarSE_0[4]]
  
  # Loading supply size parameters 
  SupplySideCoefficients_0 <-fread(paste0(project_folder, "/analysis/tablesandfigures/release/prem_setting/tables/pr_paperfull_best_guess_ra_s_wsimcost.csv"))
  SupplySideData_0 <- data.table(
    expand.grid(
      payer_id = c("M0013", "M0019", "M0035", "M0062", "M0077", "M0080"),
      year = c("2015", "2016"))
  )
  
  # Create beta_{5j}'A_j
  names(SupplySideCoefficients_0) <- c("var", "mod1", "mod2", "mod3", "mod4")
  SupplySideCoefficients_0[, year := ifelse(grepl("2015", var), "2015", NA), ]
  SupplySideCoefficients_0[, year := ifelse(grepl("2016", var), "2016", year), ]
  SupplySideCoefficients_0[, payer_id := ifelse(grepl("M0013", var), "M0013", NA), ]
  SupplySideCoefficients_0[, payer_id := ifelse(grepl("M0019", var), "M0019", payer_id), ]
  SupplySideCoefficients_0[, payer_id := ifelse(grepl("M0035", var), "M0035", payer_id), ]
  SupplySideCoefficients_0[, payer_id := ifelse(grepl("M0062", var), "M0062", payer_id), ]
  SupplySideCoefficients_0[, payer_id := ifelse(grepl("M0077", var), "M0077", payer_id), ]
  SupplySideCoefficients_0[, payer_id := ifelse(grepl("M0080", var), "M0080", payer_id), ]
  constant_0 <- as.numeric(SupplySideCoefficients_0[var == "Constant"]$mod3)
  SupplySideCoefficients2015_0 <- SupplySideCoefficients_0[year == "2015" | is.na(year)]
  SupplySideCoefficients2015_0 <- SupplySideCoefficients2015_0[ , .(
    fe_beta_5 = sum(as.numeric(mod3)) + constant_0,
    year = "2015"
  ), by = c("payer_id")]
  SupplySideCoefficients2016_0 <- SupplySideCoefficients_0[year == "2016" | is.na(year)]
  SupplySideCoefficients2016_0 <- SupplySideCoefficients2016_0[ , .(
    fe_beta_5 = sum(as.numeric(mod3)) + constant_0,
    year = "2016"
  ), by = c("payer_id")]
  SupplySideFEs_0 <- rbind(SupplySideCoefficients2015_0, SupplySideCoefficients2016_0)
  SupplySideData_0 <- merge(
    SupplySideData_0,
    SupplySideFEs_0,
    by = c("payer_id", "year"),
    all.x = T
  )
  
  # Add beta_4
  SupplySideData_0$beta_4 <- as.numeric(SupplySideCoefficients_0[var == "Medical costs"]$mod3)
  
  # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ #
  # Create parameter value matrix #
  # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ #
  
  n_sims <- 9^2 + 10^2 + 26*2 + 10^2 
  # Count all the beta parameters; from individual, small group, beta_5, and beta_4, markup
  n_pars <- length(c(unlist(NamedBetaList_0), unlist(SGNamedBetaList_0))) + nrow(SupplySideData_0) + 2 
  
  # Baseline parameters:
  base_param_vec <- c(
    unlist(NamedBetaList_0), 
    unlist(SGNamedBetaList_0), 
    SupplySideData_0$fe_beta_5, 
    SupplySideData_0$beta_4[1],
    mrkup_0)
  
  param_mat <- as.data.table(matrix(0, nrow = n_sims, ncol = n_pars))
  names(param_mat) <- c(
    paste0("im.", names(unlist(NamedBetaList_0))),
    paste0("sg.", names(unlist(SGNamedBetaList_0))),  
    paste0("al.beta5.", SupplySideData_0$payer_id, "_", SupplySideData_0$year),
    "al.beta4.constant",
    "mrkup")
  for(i in 1:n_pars){
    param_mat[, i] <- base_param_vec[i]
  }
  
  # Changing the parameters as needed
  
  ## Varying moral hazard
  param_mat$im.beta2.cons[1:81] <- rep(log(seq(.7, 1.3, .075)) + param_mat$im.beta2.cons[1], times = 9, each = 1)
  param_mat$sg.beta2.cons[1:81] <- rep(log(seq(.7, 1.3, .075)) + param_mat$sg.beta2.cons[1], times = 1, each = 9)
  
  ## Varying risk aversion
  param_mat$im.beta3.cons[82:181] <- rep(seq(-.4, 1.4, .2)+param_mat$im.beta3.cons[1], times=10, each=1)
  param_mat$sg.beta3.cons[82:181] <- rep(seq(-1.8, 0, .2)+param_mat$sg.beta3.cons[1], times=1, each=10)

  ## Varying markup 
  param_mat$mrkup[182:233] <- rep(seq(0, .25, .01), times=2, each=1)
}
# ~ # ~ # ^ CREATING PARAM MAT HERE ^ # ~ # ~ #

# Create log file for outputs
sink(paste0(output_folder, "/counterfactual_log.txt"))

# Run base parameters to find households that are not used
load_base_data <- TRUE
NamedBetaList <- NamedBetaList_0
SGNamedBetaList <- SGNamedBetaList_0
SupplySideData <- SupplySideData_0
mrkup <- mrkup_0
source(paste0(counterfactual_folder, "/library/DA03bRunCounterfactualGetdataready.R")) 

ExplodedData_SG_0 <- ExplodedData_SG
ExplodedData_SG_small_0 <- ExplodedData_SG_small
ExplodedData_Ind_0 <- ExplodedData_Ind
ExplodedData_Ind_small_0 <- ExplodedData_Ind_small

# Version that only keeps rows where alpha-psi > 0; this is for cs_1 calculations
ExplodedData_SG_01 <- ExplodedData_SG_0[ExplodedData_SG_0$alpha-ExplodedData_SG_0$psi > 0, ]
ExplodedData_SG_small_01 <- ExplodedData_SG_small_0[ExplodedData_SG_small_0$alpha-ExplodedData_SG_small_0$psi > 0, ]
ExplodedData_Ind_01 <- ExplodedData_Ind_0[ExplodedData_Ind_0$alpha-ExplodedData_Ind_0$psi > 0, ]
ExplodedData_Ind_small_01 <- ExplodedData_Ind_small_0[ExplodedData_Ind_small_0$alpha-ExplodedData_Ind_small_0$psi > 0, ]

# Only keep rows in dataset where alpha - psi > .05
ExplodedData_SG_0 <- ExplodedData_SG_0[ExplodedData_SG_0$alpha-ExplodedData_SG_0$psi > alphampsicutoff, ]
ExplodedData_SG_small_0 <- ExplodedData_SG_small_0[ExplodedData_SG_small_0$alpha-ExplodedData_SG_small_0$psi > alphampsicutoff, ]
ExplodedData_Ind_0 <- ExplodedData_Ind_0[ExplodedData_Ind_0$alpha-ExplodedData_Ind_0$psi > alphampsicutoff, ]
ExplodedData_Ind_small_0 <- ExplodedData_Ind_small_0[ExplodedData_Ind_small_0$alpha-ExplodedData_Ind_small_0$psi > alphampsicutoff, ]

# Take plan_id and hhid columns from baseline data, so we know which households to keep
ExplodedData_SG_hh <- select(ExplodedData_SG_0, "plan_id", "hhid")
ExplodedData_SG_small_hh <- select(ExplodedData_SG_small_0, "plan_id", "hhid")
ExplodedData_Ind_hh <- select(ExplodedData_Ind_0, "plan_id", "hhid")
ExplodedData_Ind_small_hh <- select(ExplodedData_Ind_small_0, "plan_id", "hhid")

ExplodedData_SG_hh1 <- select(ExplodedData_SG_01, "plan_id", "hhid")
ExplodedData_SG_small_hh1 <- select(ExplodedData_SG_small_01, "plan_id", "hhid")
ExplodedData_Ind_hh1 <- select(ExplodedData_Ind_01, "plan_id", "hhid")
ExplodedData_Ind_small_hh1 <- select(ExplodedData_Ind_small_01, "plan_id", "hhid")

rm(ExplodedData_SG) 
rm(ExplodedData_SG_small)
rm(ExplodedData_Ind)
rm(ExplodedData_Ind_small)

# ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ #
# For each row of the parameter value matrix, run a counterfactual  #
# ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ #

# We only keep results for the base and the extended coverage with mandated insurance counterfactuals:
results_list_base <- list()
results_list_5 <- list()

# Make lambda shifter vec, then change to alphashifters to multiply directly on each alpha
lambdashifter_sg <- rep(seq(.73, 1, .03), times = 10, each = 1)
lambdashifter_ind <- rep(seq(.73, 1, .03), times = 1, each = 10)

## Initialize alpha shifter list
alphashifter_sg <- rep(1, 100) 
alphashifter_ind <- rep(1, 100)
for(i in 1:length(lambdashifter_sg)){
  alphashifter_sg[i] = 1/lambdashifter_sg[i]
  alphashifter_ind[i] = 1/lambdashifter_ind[i]
}

# Estimate counterfactual equil for each simulation
for(i in 1:nrow(param_mat)){
  print(paste0("ESTIMATING COUNTERFACTUAL", i))
  
  # Unpack the matrix into NamedBetaList SGNamedBetaList SupplySidedata and mrkup 
  param_vec <- as.vector(t(param_mat[i, ]))
  names(param_vec) <- names(param_mat)
  type_param_vec <- substr(names(param_mat), 1, 8) 
  
  # NamedBetaList
  NamedBetaList <- list(
    beta1 = param_vec[which(type_param_vec %in% "im.beta1")],
    beta2 = param_vec[which(type_param_vec %in% "im.beta2")],
    beta3 = param_vec[which(type_param_vec %in% "im.beta3")],
    beta0 = param_vec[which(type_param_vec %in% "im.beta0")]
  )
  
  names(NamedBetaList[[1]]) <- find_names(names(param_vec), "im.beta1")
  names(NamedBetaList[[2]]) <- find_names(names(param_vec), "im.beta2")
  names(NamedBetaList[[3]]) <- find_names(names(param_vec), "im.beta3")
  names(NamedBetaList[[4]]) <- find_names(names(param_vec), "im.beta0")
  
  #SGNamedBetaList
  SGNamedBetaList <- list(
    beta1 = param_vec[which(type_param_vec %in% "sg.beta1")],
    beta2 = param_vec[which(type_param_vec %in% "sg.beta2")],
    beta3 = param_vec[which(type_param_vec %in% "sg.beta3")],
    beta0 = param_vec[which(type_param_vec %in% "sg.beta0")]
  )
  names(SGNamedBetaList[[1]]) <- find_names(names(param_vec), "sg.beta1") 
  names(SGNamedBetaList[[2]]) <- find_names(names(param_vec), "sg.beta2")
  names(SGNamedBetaList[[3]]) <- find_names(names(param_vec), "sg.beta3")
  names(SGNamedBetaList[[4]]) <- find_names(names(param_vec), "sg.beta0")
  
  # SupplySidedata
  SupplySideData <- data.table(
    payer_id = substr(find_names(names(param_vec), "al.beta5"), 1, 5), 
    year = str_sub(find_names(names(param_vec), "al.beta5"), -4, -1),
    fe_beta_5 = param_vec[which(type_param_vec %in% "al.beta5")], 
    beta_4 = param_vec[which(type_param_vec %in% "al.beta4")]
  )
  
  # Markup
  mrkup <- param_vec["mrkup"]
  
  # Create the relevant data
  if(i == 1){
    load_base_data <- TRUE
  } else {
    load_base_data <- FALSE
  }
  source(paste0(counterfactual_folder, "/library/DA03bRunCounterfactualGetdataready.R")) 
  
  # Varying alpha through lambda multiplier
  if(i >=234  & i <= 333) {  
    ExplodedData_SG[, in_orig_cs1 := ifelse((plan_id %in% ExplodedData_SG_hh1$plan_id) & (hhid %in% ExplodedData_SG_hh1$hhid), 1, 0)]
    ExplodedData_Ind[, in_orig_cs1 := ifelse((plan_id %in% ExplodedData_Ind_hh1$plan_id) & (hhid %in% ExplodedData_Ind_hh1$hhid), 1, 0)]

    ExplodedData_SG[, in_orig_cs2 := ifelse((plan_id %in% ExplodedData_SG_hh$plan_id) & (hhid %in% ExplodedData_SG_hh$hhid), 1, 0)]
    ExplodedData_Ind[, in_orig_cs2 := ifelse((plan_id %in% ExplodedData_Ind_hh$plan_id) & (hhid %in% ExplodedData_Ind_hh$hhid), 1, 0)]

    ExplodedData_SG[, alpha :=  alpha*alphashifter_sg[i-233]]
    ExplodedData_SG_small <- ExplodedData_SG[, .(
      orig_subs_id,
      year,
      hhid = -hhid,
      plan_id,
      best_guess_AV_old,
      best_guess_sumrate,
      best_guess_netprem_old,
      grossprem_old,
      sum_concurrent_risk,
      ratingfactor = RatioSum,
      subs = ifelse(plan_id == "Outside_Option", 0, subsidyshare),
      x_av,
      alpha,
      omega,
      psi,
      fixedeffect_beta0,
      kappa = ifelse(plan_id == "Outside_Option", 1, plan_AV /(x_av * 100)),
      in_orig_cs1,
      in_orig_cs2),]

    ExplodedData_Ind[, alpha := alpha*alphashifter_ind[i-233]]
    ExplodedData_Ind_small <- ExplodedData_Ind[, .(
      orig_subs_id,
      year,
      hhid,
      plan_id,
      best_guess_AV_old,
      best_guess_sumrate,
      best_guess_netprem_old,
      grossprem_old,
      sum_concurrent_risk,
      ratingfactor = RatioSum,
      subs = ifelse(plan_id == "Outside_Option", 0, subsidyshare),
      x_av,
      alpha,
      omega,
      psi,
      fixedeffect_beta0,
      kappa =  ifelse(plan_id == "Outside_Option", 1, plan_AV /(x_av * 100)),
      in_orig_cs1,
      in_orig_cs2),]

     tryCatch(
      source(paste0(counterfactual_folder, "/specs/", spectouse, "/loop/DA03dCounterfactualCode_change_alphaDL.R")),
      error = function(e){
        resbase <- NA
        res5 <- NA
      }
    )
    save(resbase, file = paste0(output_folder, "/resbase", i, ".Rdata"))
    save(res5, file = paste0(output_folder, "/res5", i, ".Rdata"))
    
  } else if(i >=82  & i <= 181) {
    tryCatch(
      source(paste0(counterfactual_folder, "/specs/", spectouse, "/loop/DA03dCounterfactualCode_change_psiDL.R")),
      error = function(e){
        resbase <- NA
        res5 <- NA
      }
    )
    save(resbase, file = paste0(output_folder, "/resbase", i, ".Rdata"))
    save(res5, file = paste0(output_folder, "/res5", i, ".Rdata"))
    
  } else if(i >= 208 & i <= 233) {
    mrkup_0 <- .25
    tryCatch(
      source(paste0(counterfactual_folder, "/specs/", spectouse, "/loop/DA03dCounterfactualCodeDL.R")),
      error = function(e){
        resbase <- NA
        res5 <- NA
      }
    )
    save(resbase, file = paste0(output_folder, "/resbase", i, ".Rdata"))
    save(res5, file = paste0(output_folder, "/res5", i, ".Rdata"))
    mrkup_0 <- 0
    
  } else {
    tryCatch(
      source(paste0(counterfactual_folder, "/specs/", spectouse, "/DA03dCounterfactualCode.R")),
      error = function(e){
        resbase <- NA
        res5 <- NA
      }
    )
  } 
  
  # Record outcomes as a vector in param_mat
  results_list_base[[i]] <- resbase
  results_list_5[[i]] <- res5
  
  # Kondo!
  rm(ExplodedData_SG) 
  rm(ExplodedData_SG_small)
  rm(ExplodedData_Ind)
  rm(ExplodedData_Ind_small)
}

sink()

# Saving all the results from all of the counterfactual simulations
save(results_list_base, file = paste0(output_folder, "/output_base_results_test.Rdata"))
save(results_list_5, file = paste0(output_folder, "/output_5_results_test.Rdata"))
save(param_mat, file = paste0(output_folder, "/output_param_mat_test.Rdata"))
